{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append(\"..\")\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RNN:\n",
    "    def __init__(self, Wx, Wh, b):\n",
    "        self.params = [Wx, Wh, b]\n",
    "        self.grads = [np.zeros_like(Wx), np.zeros_like(Wh), np.zeros_like(b)]\n",
    "        self.cache = None\n",
    "    \n",
    "    def forward(self, x, h_prev):\n",
    "        Wx, Wh, b = self.params\n",
    "        t = np.dot(h_prev, Wh) + np.dot(x, Wx) + b\n",
    "        h_next = np.tanh(t)\n",
    "        \n",
    "        self.cache = (x, h_prev, h_next)\n",
    "        return h_next\n",
    "    \n",
    "    def backward(self, dh_next):\n",
    "        Wx, Wh, b = self.params\n",
    "        x, h_prev, h_next = self.cache\n",
    "        \n",
    "        dt = dh_next * (1 - h_next ** 2)\n",
    "        db = np.sum(dt, axis=0)\n",
    "        dWh = np.dot(h_prev.T, dt)\n",
    "        dh_prev = np.dot(dt, Wh.T)\n",
    "        dWx = np.dot(x.T, dt)\n",
    "        dx = np.dot(dt, Wx.T)\n",
    "        \n",
    "        self.grads[0][...] = dWx\n",
    "        self.grads[1][...] = dWh\n",
    "        self.grads[2][...] = db\n",
    "        \n",
    "        return dx, dh_prev"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TimeRNN:\n",
    "    def __init__(self, Wx, Wh, b, stateful=False):\n",
    "        self.params = [Wx, Wh, b]\n",
    "        self.grads = [np.zeros_like(Wx), np.zeros_like(Wh), np.zeros_like(b)]\n",
    "        self.layers = None\n",
    "        \n",
    "        self.h, self.dh = None, None\n",
    "        self.stateful = stateful\n",
    "    \n",
    "    def set_state(self, h):\n",
    "        self.h = h\n",
    "    \n",
    "    def reset_state(self):\n",
    "        self.h = None\n",
    "    \n",
    "    def forward(self, xs):\n",
    "        Wx, Wh, b = self.params\n",
    "        N, T, D = xs.shape\n",
    "        D, H = Wx.shape\n",
    "        \n",
    "        self.layers = []\n",
    "        hs = np.empty((N, T, H), dtype='f')\n",
    "        \n",
    "        if not self.stateful or self.h is None:\n",
    "            self.h = np.zeros((N, H), dtype='f')\n",
    "        \n",
    "        for t in range(T):\n",
    "            layer = RNN(*self.params)\n",
    "            self.h = layer.forward(xs[:, t, :], self.h)\n",
    "            hs[:, t, :] = self.h\n",
    "            self.layers.append(layer)\n",
    "        \n",
    "        return hs\n",
    "    \n",
    "    def backward(self, dhs):\n",
    "        Wx, Wh, b = self.params\n",
    "        N, T, H = dhs.shape\n",
    "        D, H = Wx.shape\n",
    "        \n",
    "        dxs = np.empty((N, T, D), dtype='f')\n",
    "        dh = 0\n",
    "        grads = [0, 0, 0]\n",
    "        for t in reversed(range(T)):\n",
    "            layer = self.layers[t]\n",
    "            dx, dh = layer.backward(dhs[:, t, :] + dh)\n",
    "            dxs[:, t, :] = dx\n",
    "            \n",
    "            for i, grad in enumerate(layer.grads):\n",
    "                grads[i] += grad\n",
    "        \n",
    "        for i, grad in enumerate(grads):\n",
    "            self.grads[i][...] = grad\n",
    "        self.dh = dh\n",
    "        \n",
    "        return dxs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from common.time_layers import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleRnnlm:\n",
    "    def __init__(self, vocab_size, wordvec_size, hidden_size):\n",
    "        V, D, H = vocab_size, wordvec_size, hidden_size\n",
    "        rn = np.random.randn\n",
    "        \n",
    "        # initialize weight\n",
    "        embed_W = (rn(V, D) / 100).astype('f')\n",
    "        rnn_Wx = (rn(D, H) / np.sqrt(D)).astype('f')\n",
    "        rnn_Wh = (rn(H, H) / np.sqrt(H)).astype('f')\n",
    "        rnn_b = np.zeros(H).astype('f')\n",
    "        affine_W = (rn(H, V) / np.sqrt(H)).astype('f')\n",
    "        affine_b = np.zeros(V).astype('f')\n",
    "\n",
    "        # create layer\n",
    "        self.layers = [\n",
    "            TimeEmbedding(embed_W),\n",
    "            TimeRNN(rnn_Wx, rnn_Wh, rnn_b, stateful=True),\n",
    "            TimeAffine(affine_W, affine_b)\n",
    "        ]\n",
    "        self.loss_layer = TimeSoftmaxWithLoss()\n",
    "        self.rnn_layer = self.layers[1]\n",
    "\n",
    "        # gather all weight and gradients\n",
    "        self.params, self.grads = [], []\n",
    "        for layer in self.layers:\n",
    "            self.params += layer.params\n",
    "            self.grads += layer.grads\n",
    "    \n",
    "    def forward(self, xs, ts):\n",
    "        for layer in self.layers:\n",
    "            xs = layer.forward(xs)\n",
    "        loss = self.loss_layer.forward(xs, ts)\n",
    "        return loss\n",
    "    \n",
    "    def backward(self, dout=1):\n",
    "        dout = self.loss_layer.backward(dout)\n",
    "        for layer in reversed(self.layers):\n",
    "            dout = layer.backward(dout)\n",
    "        return dout\n",
    "    \n",
    "    def reset_state(self):\n",
    "        self.rnn_layer.reset_state()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from common.optimizer import SGD\n",
    "from dataset import ptb\n",
    "from simple_rnnlm import SimpleRnnlm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# setting of hyper-parameter\n",
    "batch_size = 10\n",
    "wordvec_size = 100\n",
    "hidden_size = 100\n",
    "time_size = 5\n",
    "lr = 0.1\n",
    "max_epoch = 100\n",
    "\n",
    "# load learning data\n",
    "corpus, word_to_id, id_to_word = ptb.load_data(\"train\")\n",
    "corpus_size = 1000\n",
    "corpus = corpus[:corpus_size]\n",
    "vocab_size = int(max(corpus) + 1)\n",
    "\n",
    "xs = corpus[:-1]\n",
    "ts = corpus[1:]\n",
    "data_size = len(xs)\n",
    "print(f\"corpus size: {corpus_size}, vocabulary size: {vocab_size}\")\n",
    "\n",
    "# variables for learning\n",
    "max_iters = data_size // (batch_size * time_size)\n",
    "time_idx = 0\n",
    "total_loss = 0\n",
    "loss_count = 0\n",
    "ppl_list = []\n",
    "\n",
    "# create model\n",
    "model = SimpleRnnlm(vocab_size, wordvec_size, hidden_size)\n",
    "optimizer = SGD(lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate starting position of each minibatch\n",
    "jump = (corpus_size - 1) // batch_size\n",
    "offsets = [i * jump for i in range(batch_size)]\n",
    "\n",
    "for epoch in range(max_epoch):\n",
    "    for iter in range(max_iters):\n",
    "        # get minibatch\n",
    "        batch_x = np.empty((batch_size, time_size), dtype='i')\n",
    "        batch_t = np.empty((batch_size, time_size), dtype='i')\n",
    "        for t in range(time_size):\n",
    "            for i, offset in enumerate(offsets):\n",
    "                batch_x[i, t] = xs[(offset + time_idx) % data_size]\n",
    "                batch_t[i, t] = ts[(offset + time_idx) % data_size]\n",
    "            time_idx += 1\n",
    "        \n",
    "        # calculate gradient and update parameters\n",
    "        loss = model.forward(batch_x, batch_t)\n",
    "        model.backward()\n",
    "        optimizer.update(model.params, model.grads)\n",
    "        total_loss += loss\n",
    "        loss_count += 1\n",
    "        \n",
    "    # evaluate perplexity every epoch\n",
    "    ppl = np.exp(total_loss / loss_count)\n",
    "    print(f\"| epoch {epoch + 1} | perplexity {ppl:.2f}\")\n",
    "    ppl_list.append(ppl)\n",
    "    total_loss, loss_count = 0, 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(range(max_epoch), ppl_list)\n",
    "plt.xlabel(\"epochs\")\n",
    "plt.ylabel(\"perplexity\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from common.trainer import RnnlmTrainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = SimpleRnnlm(vocab_size, wordvec_size, hidden_size)\n",
    "optimizer = SGD(lr)\n",
    "trainer = RnnlmTrainer(model, optimizer)\n",
    "\n",
    "trainer.fit(xs, ts, max_epoch, batch_size, time_size)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
